"""
YOLO图形检测器
"""
import cv2
import numpy as np
import torch
import os
from config.settings import DETECTION_MODE_CONFIG

class YoloDetector:
    def __init__(self):
        self.model = None
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.load_model()
    
    def load_model(self):
        """加载YOLO模型"""
        try:
            # 尝试加载YOLO模型
            model_path = DETECTION_MODE_CONFIG["yolo_model_path"]
            if os.path.exists(model_path):
                self.model = torch.hub.load('ultralytics/yolov5', 'custom', path=model_path)
            else:
                # 如果自定义模型不存在，使用预训练的yolov5s模型
                self.model = torch.hub.load('ultralytics/yolov5', 'yolov5s')
            
            self.model.to(self.device)
            self.model.eval()
        except Exception as e:
            print(f"警告：无法加载YOLO模型: {e}")
            self.model = None
    
    def detect_shapes(self, image):
        """使用YOLO检测图像中的图形"""
        if self.model is None:
            # 如果模型未加载，返回空结果
            return image, []
        
        try:
            # YOLO检测
            results = self.model(image)
            
            # 解析检测结果
            detections = results.pandas().xyxy[0]
            shapes = []
            
            # 在图像上绘制检测结果
            result_image = image.copy()
            
            for _, detection in detections.iterrows():
                # 获取边界框坐标
                x1, y1, x2, y2 = int(detection['xmin']), int(detection['ymin']), \
                                int(detection['xmax']), int(detection['ymax'])
                confidence = detection['confidence']
                class_name = detection['name']
                
                # 过滤低置信度的检测
                if confidence < 0.5:
                    continue
                
                # 绘制边界框
                cv2.rectangle(result_image, (x1, y1), (x2, y2), (0, 255, 0), 2)
                
                # 绘制标签
                label = f"{class_name}: {confidence:.2f}"
                cv2.putText(result_image, label, (x1, y1 - 10),
                           cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
                
                # 添加到形状列表
                shapes.append({
                    "name": class_name,
                    "confidence": confidence,
                    "bbox": [x1, y1, x2, y2]
                })
            
            return result_image, shapes
        except Exception as e:
            print(f"YOLO检测出错: {e}")
            return image, []
